package org.znxs.znagent_s.rag.myDefinition.githubDocumentReader;

import com.alibaba.cloud.ai.document.DocumentParser;
import lombok.extern.slf4j.Slf4j;
import org.kohsuke.github.GHContent;
import org.springframework.ai.document.Document;
import org.springframework.ai.document.DocumentReader;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.Objects;

/**
 * 自定义文档内容读取器 通过采取爬取github网页
 */
@Slf4j
public class MyGithubDocumentReader implements DocumentReader {
    private final DocumentParser parser;

    private GitHubResource gitHubResource;

    private List<GitHubResource> gitHubResourceList;


    public MyGithubDocumentReader(GitHubResource gitHubResource, DocumentParser parser) {
        this.gitHubResource = gitHubResource;
        this.parser = parser;
    }

    public MyGithubDocumentReader(List<GitHubResource> gitHubResourceList, DocumentParser parser) {
        this.gitHubResourceList = gitHubResourceList;
        this.parser = parser;
    }

    @Override
    public List<Document> get() {
        ArrayList<Document> documents = new ArrayList<>();
        if (!Objects.isNull(gitHubResourceList) && !gitHubResourceList.isEmpty()) {
            processResourceList(documents, gitHubResourceList);
        } else if (!Objects.isNull(gitHubResource)) {
            loadDocument(documents, gitHubResource);
        }

        return documents;
    }

    private void processResourceList(List<Document> documents, List<GitHubResource> gitHubResourceList) {
        // 批量加载
        for (GitHubResource gitHubResource : gitHubResourceList) {
            loadDocument(documents, gitHubResource);
        }

    }


    private void loadDocument(List<Document> documents, GitHubResource gitHubResource) {
        // 非空判断
        if (gitHubResource == null) {
            return;
        }
        // 加载github
        try {
            List<Document> documentList = parser.parse(gitHubResource.getInputStream());
            for (Document document : documentList) {
                GHContent ghContent = gitHubResource.getText();
                Map<String, Object> metadata = document.getMetadata();
                metadata.put("github_git_url", ghContent.getGitUrl());
                metadata.put("github_download_url", ghContent.getDownloadUrl());
                metadata.put("github_html_url", ghContent.getHtmlUrl());
                metadata.put("github_url", ghContent.getUrl());
                metadata.put("github_file_name", ghContent.getName());
                metadata.put("github_file_path", ghContent.getPath());
                metadata.put("github_file_sha", ghContent.getSha());
                metadata.put("github_file_size", Long.toString(ghContent.getSize()));
                metadata.put("github_file_encoding", ghContent.getEncoding());
                documents.add(document);
            }
        } catch (IOException e) {
            throw new RuntimeException("运行GitHub document reader 失败", e);
        }
    }


}
